package com.minigame.common;

import java.util.*;

import static java.util.Optional.ofNullable;
import static java.util.function.Function.identity;
import static java.util.stream.Collectors.*;

/**
 * 树结构转换工具
 */
public class TreeUtil {

    /**
     * 节点结构接口
     * 实体字段约束: id,parentId
     *
     * @param <K> id类型
     */
    public interface NodeWrapper<K> {
        K getId();
        K getParentId();
    }

    /**
     * 树结构接口
     * 实体字段约束: id,parentId,children
     *
     * @param <K> id类型
     * @param <T> 树结构类型
     */
    public interface TreeWrapper<K, T extends TreeWrapper<K, T>> extends NodeWrapper<K> {
        List<T> getChildren();
        void setChildren(List<T> children);
    }

    /**
     * 根据记录集合构建不同根节点的树结构集合
     *
     * @param nodes 节点集合
     * @param <K>   id类型
     * @param <T>   树结构类型
     * @return 所有根节点的树结构集合
     */
    public static <K, T extends TreeWrapper<K, T>> Map<K, List<T>> treesOf(List<T> nodes) {
        return sortedTreesOf(nodes, null);
    }

    /**
     * 根据记录集合构建指定根节点的树结构集合
     *
     * @param nodes  节点集合
     * @param rootId 根id
     * @param <K>    id类型
     * @param <T>    树结构类型
     * @return 指定根节点的树结构集合
     */
    public static <K, T extends TreeWrapper<K, T>> List<T> treeOf(List<T> nodes, K rootId) {
        return sortedTreeOf(nodes, rootId, null);
    }

    /**
     * 根据记录集合构建不同根节点的树结构集合, 并根据指定规则排序
     *
     * @param nodes      节点集合
     * @param comparator 排序规则
     * @param <K>        id类型
     * @param <T>        树结构类型
     * @return 所有根节点的树结构集合
     */
    public static <K, T extends TreeWrapper<K, T>> Map<K, List<T>> sortedTreesOf(List<T> nodes, Comparator<T> comparator) {
        Map<K, List<T>> childrenTrees = sortedChildrenTrees(nodes, comparator);
        retainsRoots(nodes, childrenTrees);
        return childrenTrees;
    }

    /**
     * 根据记录集合构建指定根节点的树结构集合, 并根据指定规则排序
     *
     * @param nodes      节点集合
     * @param rootId     根id
     * @param comparator 排序规则
     * @param <K>        id类型
     * @param <T>        树结构类型
     * @return 指定根节点的树结构集合
     */
    public static <K, T extends TreeWrapper<K, T>> List<T> sortedTreeOf(List<T> nodes, K rootId, Comparator<T> comparator) {
        Map<K, List<T>> childrenTrees = sortedChildrenTrees(nodes, comparator);
        return childrenTrees.get(rootId);
    }

    /**
     * 根据记录集合生成所有子树
     *
     * @param nodes 节点集合
     * @param <K>   id类型
     * @param <T>   树结构类型
     * @return 所有根节点的树结构集合
     */
    public static <K, T extends TreeWrapper<K, T>> Map<K, List<T>> childrenTrees(List<T> nodes) {
        return sortedChildrenTrees(nodes, null);
    }

    /**
     * 1.排序
     * 2.按 parentId 分组
     * 3.parent 与 children 关联
     *
     * @param nodes      记录集合
     * @param comparator 排序规则 [空则顺序不变]
     * @param <K>        id类型
     * @param <T>        树结构类型
     * @return key 为 parentId, value 为其子集
     */
    public static <K, T extends TreeWrapper<K, T>> Map<K, List<T>> sortedChildrenTrees(List<T> nodes, Comparator<T> comparator) {
        Map<K, List<T>> childrenMap = nodes.stream().sorted(ofNullable(comparator).orElse((t1, t2) -> 0)).collect(groupingBy(T::getParentId));
        nodes.forEach(node -> ofNullable(childrenMap.get(node.getId())).ifPresent(node::setChildren));
        return childrenMap;
    }

    /**
     * 根据当前节点id返回节点所在路径
     * @param nodes 节点记录集合
     * @param id 当前节点id
     * @param includeCurrent 是否包含当前节点
     * @param <K> id类型
     * @param <N> 节点类型
     * @return 节点所在路径
     */
    public static <K, N extends NodeWrapper<K>> List<N> pathOf(List<N> nodes, K id, boolean includeCurrent) {
        List<N> path = new ArrayList<>();
        Map<K, N> nodeMap = nodes.stream().collect(toMap(N::getId, identity()));
        N current = nodeMap.get(id);
        if (current == null) {
            return path;
        }
        if (includeCurrent) {
            path.add(current);
        }
        while ((current = nodeMap.get(current.getParentId())) != null) {
            path.add(0, current);
        }
        return path;
    }

    /**
     * 删除分支节点, 保留根节点
     *
     * @param nodes       节点集合
     * @param childrenMap key 为 parentId, value 为其子集
     * @param <K>         id类型
     * @param <T>         树结构类型
     */
    private static <K, T extends TreeWrapper<K, T>> void retainsRoots(List<T> nodes, Map<K, List<T>> childrenMap) {
        Set<K> ids = nodes.stream().map(T::getId).collect(toSet());
        childrenMap.keySet().removeIf(ids::contains);
    }
}
